qwen35: DSpark speculative decoding (Qwen3.8-27B drafters) - #625
qwen35: DSpark speculative decoding (Qwen3.8-27B drafters)#625davide221 wants to merge 19 commits into
Conversation
Wire the DSpark drafter heads (low-rank Markov bigram correction + confidence head) into the qwen35 spec-decode loop, so Qwen3.8-27B DSpark drafters (e.g. RadixArk/Qwen3.8-27B-DSpark) run with full head support: - spec loop: markov-corrected greedy chain (fused single-graph variant with non-fused fallback) replaces plain argmax projection when the drafter ships DSpark heads; DDTree candidate top-k gets the markov bias too. Env-gated: DFLASH_QWEN35_DSPARK, DFLASH_QWEN35_FUSED_DSPARK, DFLASH_QWEN35_DSPARK_TREE (all default on). - target capture layers now follow the drafter GGUF's dflash.target_layer_ids instead of the evenly-spaced derivation; the Qwen3.8 drafter is trained on layers 4/16/28/40/52, not 1/16/31/46/61. - draft loader: dflash.mask_token_id from the drafter GGUF wins over the family default (Qwen3.8 drafter uses 248077, default was 248070), and optional YaRN rope scaling keys are parsed into DraftWeights. - draft graph: rope calls honor the drafter's YaRN config (previously hardcoded plain NEOX rope). - Qwen35DFlashTarget exposes lm_head for the fused head path. - convert_dflash_to_gguf.py: handle single-file DSpark releases (markov/ confidence heads inline in model.safetensors), transformers>=5 nested rope_parameters and dflash_config.mask_token_id, and emit YaRN scaling metadata. The confidence-gate adaptive block length is not wired yet (q_len sizes the per-request step buffers); the chain runs with the gate off.
Verify/accept now run over v_len (the drafted chain's actual length) instead of the buffer-sizing q_len, so the DSpark confidence gate's adaptive block truncation is structurally supported. The gate itself stays off by default (DFLASH_QWEN35_DSPARK_CONFIDENCE_THRESHOLD=0): with the RadixArk Qwen3.8 drafter, any threshold in 0.1-0.5 truncates to the same short chain regardless of value, so the confidence scores coming out of the shared head path look mis-scaled and need a separate investigation before the gate can help. threshold=0 is bench-verified regression-free.
- ggml_ssm_conv_step: one kernel for the causal-conv decode/verify step (history window + silu(conv) + in-place history write-back + optional rollback window copy) replacing transpose/concat/ssm_conv/silu/cpy. - ggml_gated_delta_net_set_raw_gates: the GDN kernel applies sigmoid(beta) and softplus(alpha + dt_bias) * A itself. - ADD + RMS_NORM + MUL fusion (residual add materialized alongside the normalized output) in the CUDA/HIP graph evaluator. - legacy pool MAX_BUFFERS 256 -> 1024: LUCE_Q8_MEMO holds ~300 pooled buffers per evaluation; a full pool freed in-flight buffers with cudaFree and produced illegal memory accesses on long prefills.
Rename the RDNA small-tile macro to GGML_CUDA_MMQ_SMALL_TILE and apply it to IQ4_XS/Q4_K/Q5_K/Q6_K/Q8_0 in addition to the ROCmFPX formats. At spec-decode verify widths (N<=16) the 128-row tile leaves a 5120-row projection with only 40 blocks on a 64-CU gfx1201; 64x64/4-warp tiles measured +12-23% on those shapes (verify step 43.8 -> 39.7 ms on Qwen3.8-27B) at ~8% prefill cost.
- loader places attn_gate|attn_qkv and ssm_beta|ssm_alpha back to back and exposes zero-copy stacked aliases (L.wqkv_z, L.ssm_ba): one GEMV each instead of two (DFLASH_QWEN35_NO_STACK=1 disables). - FFN uses ggml_swiglu_split so the backend fuses gate/up/GLU into one vector kernel at decode. - DeltaNet block: single l2_norm over the q|k slab, ggml_ssm_conv_step, raw-gate gated_delta_net (in place, no state copy), no q/k head repeat (the kernel broadcasts). DFLASH_QWEN35_NO_FUSED_KERNELS=1 keeps the op-by-op graph for A/B. - DFLASH_KV_ROTATE=0 skips the FWHT K/Q rotation (precision-neutral with q8_0/f16 caches, two fewer launches per attention layer). Qwen3.8-27B IQ4_XS on R9700: plain decode 30.4 -> 33.8 tok/s with identical greedy output.
- Qwen35AdaptiveSpecPolicy: EMA of accepted draft tokens per step; below 0.8*(spec_step_ratio-1) the loop runs a burst of plain-decode steps (seed-only verify, no drafter/heads/snapshot/rollback, features still captured) and probes again afterwards. Env DFLASH_QWEN35_SPEC_STEP_RATIO (default 1.7, 0 disables) and DFLASH_QWEN35_AR_BURST (default 40). Low-acceptance prose 28.1 -> 32.4 tok/s, code/mixed unchanged. - Confidence gate now uses the fused Markov graph and truncates on the host; DFLASH_QWEN35_DSPARK_CONF_DEBUG=1 prints per-position scores. - spec-profile hooks for the chain path (project/snapshot/verify/ rollback/feature).
launch_fattn was told the vec kernel consumes D keys per step; it walks nthreads (128) per step, so a 256-key window at head_dim 256 ran as one block per head. Passing nthreads lets it use two blocks per head plus the combine pass: Qwen3.8-27B plain decode 34.3 -> 34.6 tok/s on R9700, identical output.
The first spec step after a plain-decode burst updates the acceptance EMA with alpha 0.5 so a stream that became predictable leaves plain decode immediately; step ratio and start value keep the measured best balance (45.7 / 31.8 / 40.4 tok/s code / prose / mixed).
The break-even acceptance now follows live EMAs of the spec-step and plain-step wall times (default 1.9 until both are measured), so it is right for any drafter block size (width-8 DSpark and width-16 DFlash measure ~1.8 on gfx1201).
DFlash 2 (z-lab/inco, e.g. z-lab/Qwen3.8-27B-DFlash2) is the DFlash backbone plus a grouped dynamic causal conv around attention and MLP in every layer and a candidate selector head (top-k lm_head candidates per block position, one path scored by a low-rank bigram form). - converter: maps attention_conv/mlp_conv (base kernels F32, kernel projections) and candidate_selector tensors, emits dflash2.* metadata, reads block_size from dflash_config, emits SWA pattern for drafters with causal sliding layers. - loader: DraftConvWeights per layer, DraftSelectorWeights, shape checks. - draft graph: conv prepare/finish (two taps over the block, per-element base + per-group dynamic coefficient) in both the stateless and the cached-KV builders. - selector chain: top-k via the target's GPU top-k (kMaxK 8 -> 16), one cached graph for hproj + codebook row gathers, host path search. - spec loop uses the selector before the DSpark/argmax paths. Qwen3.8-27B IQ4_XS on R9700, q8_0 drafter, greedy: 109.9 code / 50.7 prose / 111.8 mixed tok/s (DSpark drafter: 45.6 / 32.4 / 38.6); avg 5.9-6.0 accepted tokens per 8-token block on code, ~2.7 on prose.
With the 64-row/4-warp tile the mmq_x=32 instantiation runs at 180 GB/s on gfx1201 (17408x5120 IQ4_XS) against 443 GB/s at mmq_x=16 and 315 at 48, so N=17..32 batches (DDTree budgets, prefill remainders) took 2.4x longer than N=16 or N=40. Choose the next tile instead.
Resolves the conflicts with SpecLA (#606) and the packed concurrent prefill work (#595): - SSM_CONV op_params[0]: 1 stays the SpecLA heavy-light conv, the dflash fused conv step now uses 2; the CUDA dispatcher and supports_op handle both. - GDN raw-gate mode no longer uses src[8]/op_params[2] (compact-decode slot ids and the SpecLA marker): the loader builds one f32 [dt_bias | A] tensor per DeltaNet layer (TargetLayer::ssm_gate_ba, own small buffer) and ggml_gated_delta_net_set_raw_gates() attaches it as src[9] with op_params[10] = 1. - gated_delta_net.cu launchers carry both the active-slot arguments and the raw-gate parameters. - build_delta_net_block keeps main's token-axis segment structure; the stacked (z|qkv) and (beta|alpha) projections, the fused conv step, raw gates, single l2_norm over q|k and the head-broadcast shortcut apply on the plain single-sequence chain path only (ragged, compact-decode, SpecLA and chunked paths take main's materialized ops). - MMQ tiles: Q4_K keeps main's 128x64 (LUCEBOX_RDNA_MMQ_Y); IQ4_XS, Q5_K, Q6_K, Q8_0 and the ROCmFPX formats keep the 64x64 small tile. - Spec loop: DSpark Markov top-k stays available inside the non-conditional DDTree branch; SpecLA's conditional-draft path and draft-KV flag are preserved. Verified on lucebox8 (R9700): AR 34.6 tok/s, DFlash2 110/54/115 (code/prose/mixed), DSpark 48/32/40.
|
Bots been following this and testing it. They have noticed a regression TP/NVLink speed has dropped significantly vs old DFlash. They have a fix, but still cleaning and testing. |
The DDTree verify path fell back to the generic per-token GDN kernel (61-196 us/layer on gfx1201) because the grouped-cols kernel had no parent_ids handling. Port the DFS branch-transition state reload into the grouped kernel: at parent_ids[t] != t-1 the register state shard reloads from the parent's stored intermediate state (same-thread read-after-write, no barrier), root-level siblings reset to the pre-block state, and intermediates are written in tree mode so later branches can read them. Verified numerically against the generic tree kernel on a 13-node branchy tree (max rel diff 8.4e-7, reduction-order noise only); end-to-end DDTree budget-12 on the R9700 matches text output on like-for-like runs at +2% tok/s.
DDTree branches were chosen by raw per-position top-k log-probs, ignoring the DFlash2 selector entirely (it only improved the chain path). Factor the chain selector into dflash2_score_candidates() + a host-side branch-conditioned topk, and feed DDTree through build_ddtree_conditional: each expansion scores candidates as logp + selector compatibility with the branch's actual parent, log-softmax-normalized per position so cumulative best-first comparisons across depths stay on a log-prob scale (without the normalization the raw dot term mis-allocates the budget: code 126 -> 106). DFLASH_QWEN35_DSPARK_TREE/raw top-k remain the fallback; DFLASH_QWEN35_DFLASH2_TREE=0 disables. R9700, budget 12 (code/prose/mixed tok/s): raw tree 126/56/112, selector tree 121/63/116, chain 112/62/124. The selector tree no longer collapses on low-acceptance content; chain remains the serving default.
The DFLASH27B_CHUNKED path was unusable on ROCm: a 512-token prefill forward took 39 s. Two causes, both fixed: - the [CS, CS] triangular solve at CS = 64 missed ggml-cuda's fast warp kernel (k <= 32) and fell into cublasStrsmBatched; CS = 32 keeps the solve on the fast path (graph cap raised to 32k nodes to match), - the per-chunk slices are strided views, which pushed every chunk matmul into cublasGemmBatchedEx; on ROCm that API stages its device pointer arrays through per-call pinned host allocations (~1 ms of hipHostMalloc/hipFree per node). ggml_cont on the sliced operands restores the strided-batched fast path. Result: 39 s -> 0.6 s per 512-token forward, output verified ~1e-6 against the sequential kernel at T = 64..2048 including padded chunks. Still OFF by default: on gfx1201 the sequential fused GDN kernel wins (514 ms vs 667 ms per forward; the ~20k-node chunk graph costs more in launches than it saves in serialization). DFLASH27B_CHUNKED=1 opts in, and the gate is per-call now, so enabling it no longer disables the raw-gate fusion on the decode path as a side effect. Also: env-gated DFLASH_PREFILL_TIMING=1 build/alloc/compute breakdown per prefill ubatch, and drop the ROCMFP requant experiment script that slipped into the merge commit (the format was refuted for this target). R9700 regression check (pure-IQ4_XS target): AR 36.4-36.6, DFlash2 spec 111/62/123 code/prose/mixed, prefill 1036/1102/1038 tok/s @512/2048/6000.
The fp64 RoPE path (required for Qwen3.5-family freq_base=1e7, see the fp32 precision wall note) computed pow(double, double) per element. On RDNA4 that libcall made rope_multi the second-largest prefill kernel: 692 us per launch at n_tokens=512 vs 76 us for the fp32 upstream kernel, ~33 ms of a 514 ms 512-token prefill forward. Replace pow() with binary exponentiation (<= 7 double multiplies for exponent < 128), keeping the large-freq_base precision to within 1 ulp. R9700: 512-token prefill forward 514 -> 414-423 ms (prefill ~996 -> ~1225 tok/s); DFlash2 spec decode 112.7/62.2/125.1 code/prose/mixed (from 111/61.5/123.4) with per-position acceptance identical.
The dense hybrid types compile their MMQ instances with the 64x64 small tile (GGML_CUDA_MMQ_SMALL_TILE), which wins 12-23% at spec-decode verify widths but re-streams the weights through narrow x-tiles at prefill widths (measured +16-18% kernel time at N=512 vs the 128x128 upstream shape). The tile shape is baked into every mmq.cuh constexpr via macros, so one TU can only hold one shape. Add big-tile twin instances for IQ4_XS/Q4_K/Q5_K/Q6_K/Q8_0 that re-include mmq.cuh inside namespace lucebox_mmq_big with no tile macro, giving the 128x128 shape distinct symbols, plus bridge functions and a runtime dispatch: RDNA4 + ncols_dst >= 256 takes the big tile (measured crossover: small wins to N=64, tie at 128, big wins 16-18% at 512); everything else keeps today's path. LUCE_MMQ_BIG_PREFILL=0 disables. gfx1151 behavior unchanged. R9700, Qwen3.8-27B pure-IQ4_XS: 512-token prefill forward 414 -> 365-374 ms (prefill ~1225 -> ~1385 tok/s, past upstream llama.cpp's 1366 pp512); generated output hash-identical; spec decode unchanged at 112.6/62.4/125.4 code/prose/mixed (verify widths never take the big tile).
Plain decode streams every weight byte exactly once per token, so caching the weight stream in L2 only evicts the activations and KV that other kernels reuse. Add a nontemporal load variant (HIP sc0/sc1 bypass hints) and use it for the IQ4_XS weight words in the MMVQ vec_dot; the q8_1 activation loads keep normal caching. Scope notes from measurement (R9700): the MMVQ weight reads are wave-contiguous full cache lines, so bypassing L2 is free there (GEMV 550 -> 554 GB/s, AR decode 36.6 -> 37.0 tok/s, +1%). The same hint in the MMQ tile loader was measured 32% SLOWER (457 -> 310 GB/s at N=8: tile blocks on different CUs share cache lines, and bypassing L2 amplifies DRAM traffic), so MMQ keeps cached loads. q8_0 qs is only 2-byte aligned and keeps get_int_b2. Spec decode unchanged at 112.9/62.4/125.4; outputs identical.
|
R9700 / gfx1201 follow-up (commit Measured on Qwen3.8-27B IQ4_XS + Q8_0 DFlash2, same production flags:
On the mixed writing/reasoning/math/code/STEM/agent set, width 12 was +3.8% overall; width 16 had small regressions on a few low-acceptance prose prompts. Full HumanEval+ pass@1 was 145/164 at width 12 versus 143/164 at width 8. Across all 164 generations, aggregate decode rose 139.6 -> 178.5 tok/s (+28%) and total generation wall time fell 319.7 -> 257.1 s. Validation: gfx1201 server build passed, all 386 server/feature-gate tests passed, invalid CLI bounds are rejected, and the service launch reproduced 230.2 tok/s through the new CLI path. The flag is monolithic dense-Qwen only; capability warnings cover unsupported architectures/layer split, and remote draft use is rejected because IPC owns its graph width. |
Adds DSpark and DFlash 2 drafter support to the qwen35 family spec-decode loop, targeting the published Qwen3.8-27B checkpoints (
RadixArk/Qwen3.8-27B-DSpark: DFlash backbone + low-rank Markov bigram head + confidence head;z-lab/Qwen3.8-27B-DFlash2: DFlash backbone + grouped dynamic convs + candidate selector), plus the decode/prefill/tree optimizations listed below. Merged with currentmain(SpecLA #606, packed concurrent prefill #595); all GPU CI runners green.What's in here
Spec loop (
qwen35_backend.cpp)DFLASH_QWEN35_DSPARK,DFLASH_QWEN35_FUSED_DSPARK,DFLASH_QWEN35_DSPARK_TREE,DFLASH_QWEN35_DFLASH2_TREE.DFLASH_QWEN35_SPEC_STEP_RATIO,DFLASH_QWEN35_AR_BURST).dflash.target_layer_ids;dflash.mask_token_idfrom the GGUF wins over the family default. Either one wrong silently destroys acceptance.Draft side: DFlash 2 dynamic convs + selector in the draft graph and converter; optional YaRN rope keys parsed and honored; single-file DSpark releases convert directly.
ggml (HIP) kernels
ggml_ssm_conv_step, raw-gategated_delta_net(gate tensor[dt_bias | A]built by the loader,src[9]/op_params[10], coexists with SpecLA'ssrc[8]/op_params[2]), residual ADD+RMS_NORM+MUL fusion, in-place recurrent state. Launches per decode token 1929 -> ~1140.GGML_CUDA_MMQ_SMALL_TILE; Q4_K keeps main's 128x64), pathological mmq_x=32 tile skipped, FA vec short-KV split, legacy pool MAX_BUFFERS 256 -> 1024.DFLASH27B_CHUNKED=1opts in): chunk size 32 keeps the triangular solve on the fast warp kernel,ggml_conton chunk slices avoids ROCm's pinned-host staging incublasGemmBatchedEx(39 s -> 0.6 s per 512-token forward, verified ~1e-6). On gfx1201 the sequential fused GDN kernel still wins at ubatch 512, hence off.DFLASH_PREFILL_TIMING=1prints per-ubatch build/alloc/compute.Numbers (R9700 / gfx1201, ROCm 7.2, greedy, 300-token gens)
Target for all rows: pure-IQ4_XS requant (
llama-quantize --allow-requantize --pure ... iq4_xs, 13.85 GB body; PPL 4.359 -> 4.387, +0.65% on 40x512 chunks).Decode, tok/s (code / prose / mixed):
--ddtree --ddtree-budget 12)No single tree/chain config dominates: chain wins high-acceptance content, raw-topk tree wins code, the selector tree is the robust middle. Chain is the serving default; the tree is a flag.
Prefill: 1036 / 1102 / 1038 tok/s at 512 / 2048 / 6000 prompt tokens (hipfire published 759 / 737 / 663; upstream llama.cpp 1318 / 1277 / 1123). Long context and thinking-mode numbers in the review thread; with max_tokens below
hard_limit_reply_budget(4096) the reasoning budget hook force-closes thinking and the reply runs as plain decode.Production serving config
Greedy verification keeps the output exact whichever drafter is used. DSpark and the 3.6-DFlash drafter keep working (the z-lab Qwen3.6-27B drafter transfers to 3.8 unchanged).
Known limitations / follow-ups